/*
 * @Author: youngqj youngqj@126.com
 * @Date: 2024-11-01 22:44:39
 * @LastEditors: youngqj youngqj@126.com
 * @LastEditTime: 2024-12-04 16:28:52
 * @FilePath: /mallSystem/middleware/db_middleware.go
 * @Description:
 *
 * Copyright © Zhejiang Xiaoqu Information Technology Co., Ltd, All Rights Reserved.
 */
package middleware

import (
	"fmt"
	"xiaoquio/web/core/database"
	"xiaoquio/web/core/utils"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

// DBMiddleware 确保每个请求都能获取到数据库连接
func DBMiddleware(dbName string) gin.HandlerFunc {
	if dbName == "default" {
		dbName = "mall_system" // 默认数据库名
	}

	return func(c *gin.Context) {
		// 获取已初始化的 DBManager
		dbManager, err := database.NewDBManager("database.yaml")
		if err != nil {
			utils.ResponseNilData(c, 500, "数据库链接失败")
			c.Abort()
			return
		}

		// 从上下文中检查是否已存在数据库连接
		db, exists := c.Get("db")
		if !exists || db == nil {
			// 如果数据库连接不存在，连接指定数据库
			conn, err := dbManager.ConnectDB(dbName)
			if err != nil {

				utils.ResponseNilData(c, 500, "数据库链接失败")
				c.Abort()
				return
			}
			// 将数据库连接存储到上下文中
			c.Set(dbName, conn)
		}

		// 继续处理请求
		c.Next()
	}
}

// GetDB 从上下文中获取数据库连接
func GetDB(c *gin.Context, dbName string) (*gorm.DB, error) {
	if dbName == "default" {
		dbName = "mall_system" // 默认数据库名
	}
	db, exists := c.Get(dbName)
	if !exists || db == nil {
		return nil, fmt.Errorf("database connection not found")
	}
	return db.(*gorm.DB), nil
}
